"""
Python file contiaining the functions that I'm going to be using 
for this project.

B. Yadidya
"""

import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import timedelta
import datetime
import sys

import logging
# Set up logging to output to the terminal
logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(message)s',  # Define the log message format
    level=logging.INFO, stream=sys.stdout  # Set the logging level to INFO to capture general messages and above
)

def var_reduction(sla_var,sla,model):
    """
    sla_var - variance in sla (cm^2)
    sla - absolute sla (cm)
    model - absolute sla from model (cm)
    """
    
    return sla_var - (sla-model).var(dim='cycle_num',skipna=True)

# =====================================================================================================

# Adding some functions that I've been using for analysis and plotting - Sep 19, 2024


def explained_variance(sla, model):
    """
    Calculate the explained variance.

    sla - absolute sla (cm)
    model - absolute sla from model (cm)
    """
    total_variance = sla.var(dim='cycle_num', skipna=True)
    residual_variance = (sla - model).var(dim='cycle_num', skipna=True)
    explained_variance_value = total_variance - residual_variance
    return explained_variance_value



def coarsen_data(da, target_resolution_deg = .1):
    """
    Coarsen to a target resolution in degrees.
    
    Parameters:
        target_resolution_deg (float): Target resolution in degrees 
    
    Returns:
        xarray.DataArray: Coarsened dataset.
    """

    # Calculate coarsening factor (assuming original grid ~0.036 deg for 4 km)
    original_resolution_deg = 0.036
    factor = int(target_resolution_deg / original_resolution_deg)

    # Coarsen the data
    coarsened = da.coarsen(
        latitude=factor,
        longitude=factor,
        boundary='trim'
    ).mean()

    return coarsened
